import asyncio
import os
from aiohttp.web_fileresponse import FileResponse
from fastapi import FastAPI, staticfiles
from fastapi.responses import HTMLResponse, JSONResponse
from spacy import tokenizer
from starlette.responses import StreamingResponse
from starlette.staticfiles import StaticFiles
# from app.api.endpoints import kb_chat, chat  # 暂时注释掉以避免向量存储初始化问题
from fastapi.templating import Jinja2Templates
from fastapi import Request
from pydantic import BaseModel
from torch.cuda import device
from transformers import TextIteratorStreamer

from app.utils.user_store import create_user, verify_user, get_current_user_id  # 新增导入
from app.utils.file_processor import process_uploaded_file, is_supported_file  # 文件处理导入
from fastapi import HTTPException, status
from fastapi import UploadFile, File
from chat_with_sichuan_bot import answer_with_sichuan_bot
from app.api.endpoints import session_management  # 会话管理API导入
from app.core.memory.memory_manager import get_user_session_history  # 会话历史管理导入
app = FastAPI(
    title="LLM Application Scaffold - Separated Endpoints",
    description="An API with distinct, stateful endpoints for Knowledge Base (RAG) chat and Agent (Tool-calling) chat.",
    version="2.2.0",
)

# 注册会话管理API路由
app.include_router(session_management.router, prefix="/api/sessions", tags=["sessions"])

from pathlib import Path

# 获取项目根目录
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
templates = Jinja2Templates(directory=str(PROJECT_ROOT / "app" / "static"))

async def stream_answer_generator(question: str):
    """
    这是一个异步生成器，它会调用模型并逐个token地yield结果。
    """
    # 关键1: 创建一个 TextIteratorStreamer 实例
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    # 关键3: 遍历streamer，逐个yield出生成的文本块，并加入延时
    for new_text_chunk in streamer:
        # **关键改动**: 遍历从streamer中获取的文本块(chunk)
        for char in new_text_chunk:
            yield char  # 一次只发送一个字
            # --- 在这里控制速度！---
            # 将延时放在内层循环，确保每个字之间都有停顿
            await asyncio.sleep(0.05)

@app.get("/", response_class=HTMLResponse)
def root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


class InputData(BaseModel):
    input: str
    conversation_id: str = None
    username: str = None


# **关键改进**：异步生成器，逐字或逐块返回
async def stream_answer_generator(answer: str):
    # 模拟逐字输出（类似 ChatGPT 的效果）
    for char in answer:
        yield char
        await asyncio.sleep(0.05)  # 控制输出速度


@app.post("/api/submit")
async def receive_input(data: InputData):
    user_input = data.input
    conversation_id = data.conversation_id
    username = data.username

    # 如果提供了用户名和会话ID，使用持久化会话历史
    if username and conversation_id:
        try:
            user_id = get_current_user_id(username)
            if user_id:
                # 获取会话历史
                session_history = get_user_session_history(conversation_id, user_id)

                # 添加用户消息到历史
                session_history.add_user_message(user_input)

                # 生成回答
                answer = answer_with_sichuan_bot(user_input)

                # 添加AI回答到历史
                session_history.add_ai_message(answer)
            else:
                # 用户不存在，使用普通模式
                answer = answer_with_sichuan_bot(user_input)
        except Exception as e:
            print(f"会话历史处理出错: {e}")
            # 出错时回退到普通模式
            answer = answer_with_sichuan_bot(user_input)
    else:
        # 没有提供会话信息，使用普通模式
        answer = answer_with_sichuan_bot(user_input)

    # **关键 HTTP 头**：禁用缓冲，确保流式传输
    headers = {
        "Content-Type": "text/plain; charset=utf-8",
        "Cache-Control": "no-cache",
        "X-Accel-Buffering": "no",  # 防止 Nginx 等代理缓冲
    }

    return StreamingResponse(
        stream_answer_generator(answer),
        headers=headers,
    )
class AuthData(BaseModel):
    username: str
    password: str


@app.post("/api/register")
async def api_register(data: AuthData):
    """用户注册接口"""
    ok, msg = create_user(data.username, data.password)
    if not ok:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=msg)
    return {"message": msg}


@app.post("/api/login")
async def api_login(data: AuthData):
    """用户登录接口"""
    ok, msg = verify_user(data.username, data.password)
    if not ok:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=msg)
    return {"message": msg}


@app.get("/login", response_class=HTMLResponse)
def login_page(request: Request):
    return templates.TemplateResponse("login.html", {"request": request})


@app.get("/register", response_class=HTMLResponse)
def register_page(request: Request):
    return templates.TemplateResponse("register.html", {"request": request})


@app.post("/api/upload")
async def api_upload(file: UploadFile = File(...)):
    """接收前端上传的文件，提取文本内容并返回"""
    try:
        # 检查文件类型
        if not is_supported_file(file.filename):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"不支持的文件类型。支持的格式: PDF, DOC, DOCX"
            )

        # 读取文件内容
        content = await file.read()

        # 提取文本内容
        success, result = process_uploaded_file(file.filename, content)
        if success:
            return {
                "success": True,
                "filename": file.filename,
                "text_content": result,
                "message": f"成功提取文件内容：{file.filename}"
            }
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"文件处理失败：{result}"
            )

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"服务器处理文件时发生错误：{str(e)}"
        )
